#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <mpi.h>

#include <CL/opencl.h>
#include "clError.h"
#include "ini.h"

#include "../util.h"
#include "../common.h"
#include "helper.h"

#define ROWSIZE 2000
#define COLSIZE 4000
#define FILENAME "../2000x4000.bin"

#define INI "oclconfig.ini"
typedef struct {
   cl_device_type device_type;
} Configuration;

int ini_handler(void *data, const char *section, const char *name, const char *value) {
   Configuration *config = (Configuration *) data;
   if (strcasecmp("DEVICE_TYPE",name) == 0) {
      if (strcasecmp("CL_DEVICE_TYPE_CPU",value) == 0) {
         config->device_type = CL_DEVICE_TYPE_CPU;
      } else if (strcasecmp("CL_DEVICE_TYPE_GPU",value) == 0) {
         config->device_type = CL_DEVICE_TYPE_GPU;
      }
   }
}

/* a short function to print a message and exit */
void error_exit(char msg[]) {
    fprintf(stderr, "%s", msg);
    MPI_Finalize(); exit(EXIT_FAILURE);
}

float *calculation(int task_id, cl_command_queue queue, cl_kernel kernel, char *filename, cl_mem inAG, cl_mem inBG, cl_mem outG, int row, int col, size_t tileSize);

/* ---- main program ---- */

int main(int argc, char *argv[]) {

  int num_tasks;

  double start_time0, start_time, end_time;

  int i;
  int task_id;
  int total_task_id = 0;
  int process_num_tasks = 0;
  int total_process_num_tasks = 0;
  int verbose = 0;
  char* usage_msg = "parameters:  num_tasks [--verbose]\n";
  int required_parms = 2;
  int nprocs;
  int process_ID;
  MPI_Status status;

  if (MPI_Init(&argc, &argv) != MPI_SUCCESS) {
      fprintf(stderr, "MPI initialization error\n"); 
      exit(EXIT_FAILURE);
  }

  if (argc < required_parms) {
      error_exit(usage_msg);
  }
  num_tasks = atoi(argv[1]);
  if (num_tasks <= 0) {
      error_exit(usage_msg);
  }
  if (argc > required_parms) {
      if (strcmp(argv[required_parms], "--verbose") == 0) 
          verbose = 1;
      else
          error_exit(usage_msg);
  }

  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &process_ID);

  if (nprocs < 2) {
      error_exit("number of processes must be at least 2\n");
  }

  char hostname[128];
  int hostname_len;
  MPI_Get_processor_name(hostname,&hostname_len);
  printf("Rank %d is running on %s\n",process_ID,hostname);

  MPI_Barrier(MPI_COMM_WORLD);
  start_time0 = MPI_Wtime();

  if (process_ID == 0) {
      printf("Use input data block from \n",FILENAME);
      for (i = 0; i < num_tasks; ++i) {

          MPI_Recv(0, 0, MPI_INT, MPI_ANY_SOURCE, 0, 
                  MPI_COMM_WORLD, &status);

          task_id = i;
          MPI_Send(&task_id, 1, MPI_INT, status.MPI_SOURCE, 0, 
                  MPI_COMM_WORLD);
      }
      task_id = -1;
      for (i = 1; i < nprocs; ++i) {
          MPI_Send(&task_id, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
      }
  }
  else { // worker
      cl_device_id device_id;
      cl_context context;
      cl_command_queue queue;
      cl_kernel kernel;
      cl_int err;
      
      // Read ini configuration file
      Configuration config;
      char *home = getenv("HOME");
      if (ini_parse(INI,ini_handler,&config) < 0) {
         char* filename = (char *) malloc((strlen(home) + strlen(INI) + 2)*sizeof(char));
         strcpy(filename,home);
         strcat(filename,"/"INI);
         if (ini_parse(filename,ini_handler,&config) < 0) {
            fprintf(stderr,"Can't load config file\n");
            exit(1);
         }
         printf("Read oclconfig.ini from %s\n",home);
         free(filename);
      } else {
         printf("Read oclconfig.ini from .\n");
      }

      // Get a device 
      device_id = opencl_get_device_id(config.device_type);

      // Create a context
      context = clCreateContext(0,1,&device_id,NULL,NULL,&err);
      check_cl_error(err);

      // Create a command queue
      queue = clCreateCommandQueue(context,device_id,0,&err);
      check_cl_error(err);

      // Load and compile a kernel from source file 
      const char* sources[2] = {"pearson.cl","../common.h"};
      kernel = opencl_load_kernel_from_file(context,device_id,2,sources,"corrcoefmatrixAB_opencl");

      // Find appropriate tile size
      size_t wavefrontSize = opencl_get_wavefront_size(kernel,device_id);
      printf("Device wavefront size = %d\n",wavefrontSize);
      //cl_device_id device_id = opencl_query_device_id(queue);
      size_t workGroupSize = opencl_get_max_work_group_size(device_id);
      printf("Max work group size = %d\n",workGroupSize);
      size_t tileSize = TILESIZE;
      while (tileSize*tileSize > workGroupSize) {
         tileSize >>= 1;
      }
      printf("Tile size = %d\n",tileSize);

      cl_mem inAG, inBG, outG;
      float *out;
      int sizeIn, sizeOut;

      sizeIn = ROWSIZE*COLSIZE*sizeof(float);
      sizeOut = ROWSIZE*ROWSIZE*sizeof(float);

      inAG = clCreateBuffer(context,CL_MEM_READ_ONLY,sizeIn,NULL,&err);
      if (inAG == NULL) {
         fprintf(stderr,"device memory allocation error: %s!\n",clErrorString(err)); 
         MPI_Abort(MPI_COMM_WORLD,1);
      }
      inBG = clCreateBuffer(context,CL_MEM_READ_ONLY,sizeIn,NULL,&err);
      if (inBG == NULL) {
         fprintf(stderr,"device memory allocation error: %s!\n",clErrorString(err)); 
         MPI_Abort(MPI_COMM_WORLD,1);
      }
      outG = clCreateBuffer(context,CL_MEM_WRITE_ONLY,sizeOut,NULL,&err);
      if (outG == NULL) {
         fprintf(stderr,"device memory allocation error: %s!\n",clErrorString(err)); 
         MPI_Abort(MPI_COMM_WORLD,1);
      }

      task_id = 0;
      while (task_id >= 0) {

          MPI_Send(0, 0, MPI_INT, 0, 0, MPI_COMM_WORLD);

          MPI_Recv(&task_id, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);

          if (task_id >= 0) {
              ++process_num_tasks;
              if (verbose)
                  fprintf(stderr,"(process %d) task = %d\n", process_ID, task_id);
              start_time = MPI_Wtime();
              out = calculation(task_id,queue,kernel,FILENAME,inAG,inBG,outG,ROWSIZE,COLSIZE,tileSize);
              end_time = MPI_Wtime();
//              write_submatrix(out,ROWSIZE,ROWSIZE,ROWSIZE,ROWSIZE);
              free(out);
              fprintf(stderr,"(process %d) finished task = %d in %g\n", process_ID, task_id,(end_time-start_time));
          }
      }
      fprintf(stderr,"process %d number of tasks = %d\n", process_ID,
              process_num_tasks);

      clReleaseKernel(kernel);
      clReleaseMemObject(inAG);
      clReleaseMemObject(inBG);
      clReleaseMemObject(outG);
      clReleaseCommandQueue(queue);
      clReleaseContext(context);

  }

  MPI_Reduce(&process_num_tasks, &total_process_num_tasks, 1,
          MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);

  MPI_Barrier(MPI_COMM_WORLD);
  end_time = MPI_Wtime();

  if (process_ID == 0) {
      fprintf(stderr,"\nMPI parallel version with %d worker processes"
              ", dynamic assignment of tasks\n", 
              nprocs-1);
      fprintf(stderr,"number of tasks = %d\n", total_process_num_tasks);
      fprintf(stderr,"running time = %g\n", end_time - start_time0);
  }

  MPI_Finalize();
  return EXIT_SUCCESS;
}

float *calculation(int task_id, cl_command_queue queue, cl_kernel kernel, char *filename, cl_mem inAG, cl_mem inBG, cl_mem outG, int row, int col, size_t tileSize) {
  float *in, *out;
  int padded_row;
  int sizeIn, sizeOut;
  cl_int err;
  
  in = read_raw_matrix(filename,tileSize,row,col,&padded_row);
  if (row != padded_row) {
     fprintf(stderr,"Not multiple of tileSize\n");
     return NULL;
  } 
  sizeIn = row*col*sizeof(float);
  sizeOut = row*row*sizeof(float);
  out = (float *) malloc(sizeOut);
  
  if ((err=clEnqueueWriteBuffer(queue,inAG,CL_TRUE,0,sizeIn,in,0,NULL,NULL)) != CL_SUCCESS) {
     fprintf(stderr,"memcpy from host to device fail: %s\n",clErrorString(err));
     return NULL;
  }
  if ((err=clEnqueueWriteBuffer(queue,inBG,CL_TRUE,0,sizeIn,in,0,NULL,NULL)) != CL_SUCCESS) {
     fprintf(stderr,"memcpy from host to device fail: %s\n",clErrorString(err));
     return NULL;
  }

  err = clSetKernelArg(kernel,0,sizeof(cl_mem),&outG);
  err |= clSetKernelArg(kernel,1,sizeof(cl_mem),&inAG);
  err |= clSetKernelArg(kernel,2,sizeof(cl_mem),&inBG);
  err |= clSetKernelArg(kernel,3,sizeof(int),&padded_row);
  err |= clSetKernelArg(kernel,4,sizeof(int),&col);
  err |= clSetKernelArg(kernel,5,sizeof(int),&tileSize);
  if (err != CL_SUCCESS) {
     fprintf(stderr,"set kernel arguments fails: %s\n",clErrorString(err));
     return NULL;
  }

  size_t localSize[2] = {tileSize,tileSize};
  size_t globalSize[2] = {padded_row,padded_row};

  err = clEnqueueNDRangeKernel(queue,kernel,2,NULL,globalSize,localSize,0,NULL,NULL);
  if (err != CL_SUCCESS) {
     fprintf(stderr,"execute the kernel fails: %s\n",clErrorString(err));
     return NULL;
  }

  clFinish(queue);


  if ((err=clEnqueueReadBuffer(queue,outG,CL_TRUE,0,sizeOut,out,0,NULL,NULL)) != CL_SUCCESS) {
     fprintf(stderr,"memcpy from device to host fail: %s\n",clErrorString(err));
     return NULL;
  }

  return out;  
}



